import io
import pymysql
from gevent import monkey
import base64
import numpy as np
from flask import Flask, jsonify, request
from gevent.pywsgi import WSGIServer
from keras.models import load_model
from keras.utils import img_to_array
import tensorflow as tf
from PIL import Image

monkey.patch_all()
app = Flask(__name__)

# 定义类别索引映射字典
image_class_index = {
    "0": "苹果",
    "1": "鳄梨",
    "2": "香蕉",
    "3": "黑莓",
    "4": "蓝莓",
    "5": "西兰花",
    "6": "卷心菜",
    "7": "辣椒",
    "8": "胡萝卜",
    "9": "红辣椒",
    "10": "玉米",
    "11": "黄瓜",
    "12": "火龙果",
    "13": "茄子",
    "14": "大蒜",
    "15": "姜",
    "16": "葡萄",
    "17": "猕猴桃",
    "18": "柠檬",
    "19": "芒果",
    "20": "蘑菇",
    "21": "洋葱",
    "22": "橘子",
    "23": "花生",
    "24": "梨",
    "25": "菠萝",
    "26": "石榴",
    "27": "马铃薯",
    "28": "南瓜",
    "29": "黄豆",
    "30": "草莓",
    "31": "番茄",
    "32": "西瓜",
}

# 加载模型时设置 compile=False 避免加载损失函数报错
model = load_model("DenseNet121_No_Aug_best.h5", compile=False)
# 如果仅用于预测，此处无需调用 compile() 方法
# tf.compat.v1.keras.backend.clear_session()  可根据需要决定是否清理 session

def transform_image(img):
    """将图片调整为 224x224 并归一化"""
    image = img.resize((224, 224))
    image_array = img_to_array(image)
    image_array = image_array / 255.0  # 归一化到 [0, 1]
    test_image = np.expand_dims(image_array, axis=0)
    return test_image

def get_prediction(img):
    """获取图片预测结果，返回类别 id、类别名称和概率"""
    tensor = transform_image(img)
    pre = model.predict(tensor)
    print(pre)
    y_hat = np.argmax(pre, axis=1)[0]
    prob = str(pre[0][y_hat])
    predicted_idx = str(y_hat)
    return predicted_idx, image_class_index[predicted_idx], prob

@app.route('/predict', methods=['POST'])
def predict():
    # 从前端获取 base64 格式的图片数据
    img_base64 = request.form.get('picture')
    if not img_base64:
        return jsonify({"error": "未提供图片数据"}), 400

    try:
        image_data = base64.b64decode(img_base64)
        image = Image.open(io.BytesIO(image_data))
    except Exception as e:
        return jsonify({"error": "图片数据解析失败", "details": str(e)}), 400

    # 进行图片预测
    class_id, class_name, prob = get_prediction(image)

    # 根据预测的 class_id 查询数据库中对应的记录，获取英文名、another_name、classify、nourishment、eating、price
    cursor = None  # 初始化 cursor 为 None
    try:
        db = pymysql.connect(host="localhost", user="root", password="123456", database="fruit")
        cursor = db.cursor()
        sql = """SELECT english_name, another_name, classify, nourishment, eating, price 
                 FROM fruit_vegetable WHERE id = %s"""
        cursor.execute(sql, (class_id,))
        result = cursor.fetchone()
        if result:
            english_name, another_name, classify, nourishment, eating, price = result
        else:
            english_name = another_name = classify = nourishment = eating = price = "未找到对应记录"
    except pymysql.Error as e:
        english_name = another_name = classify = nourishment = eating = price = f"查询出错: {e}"
    finally:
        if cursor:
            cursor.close()
        if 'db' in locals():
            db.close()

    return jsonify({
        'class_id': class_id,
        'class_name': class_name,
        'prob': prob,
        'english_name': english_name,
        'another_name': another_name,
        'classify': classify,
        'nourishment': nourishment,
        'eating': eating,
        'price': price
    })

if __name__ == '__main__':
    server = WSGIServer(('0.0.0.0', 5000), app)
    server.serve_forever()
